import builtins
import os
from typing import Union
from .base_dataset import BaseDataset

class RealStateDataset(BaseDataset):
    def __init__(
        self,
        filename_ls_path: str,
        dataset_dir: str,
        random_mask_prob = 0.5,
        disp_name: str = 'DL3DV',
        input_size: int = 518,
        validate: bool = False,
        max_disparity: float = 20.0, #TODO: 什么值比较好
        mask_expand_radius: int = 5,
        **kwargs,
    ) -> None:
        super().__init__(
            filename_ls_path = filename_ls_path,
            dataset_dir = dataset_dir,
            random_mask_prob = random_mask_prob,
            disp_name = disp_name,
            input_size = input_size,
            validate = validate,
            max_disparity = max_disparity,
            mask_expand_radius = mask_expand_radius,
            **kwargs)
        
    def get_path(self, index):
        index = index % builtins.len(self.file_list)
        file_sub_path = self.file_list[index]
        video_path = os.path.join(self.dataset_dir,'dataset_video', f'{file_sub_path}.mp4')
        point_path = os.path.join(self.dataset_dir,'point_cloud_ply', file_sub_path,'pc_0.ply')
        camera_path = os.path.join(self.dataset_dir,'point_cloud_ply', file_sub_path,'camera.pt')
        mask_dir_path = os.path.join(self.dataset_dir,'mask', file_sub_path,'0')
        return {
            "video_path": video_path,
            "point_path": point_path,
            "camera_path": camera_path,
            "mask_dir_path": mask_dir_path,
        }